# mopsorunner.py

import numpy as np
import pandas as pd
from dataclasses import dataclass
import matplotlib.pyplot as plt
import random
import math, os
import copy
from tqdm import tqdm # For progress bar

# Assuming config.py and environment.py are in the same directory or accessible
from config import get_default_configs, SystemConfig, GeneratorConfig, StorageConfig
from environment import PowerSystemEnv, State

# --- MOPSO Configuration ---
@dataclass
class MOPSOConfig:
    pop_size: int = 400       # Reduced for potentially faster testing
    n_generations: int = 100  # Increased generations might be needed
    # PSO parameters
    w_max: float = 0.9        # Initial inertia weight
    w_min: float = 0.2        # Final inertia weight
    c1: float = 1.5           # Cognitive coefficient
    c2: float = 1.5           # Social coefficient
    # Archive size limit (optional, -1 for unlimited)
    archive_size_limit: int = 100
    # Velocity limits (to prevent explosion, adjust based on variable ranges)
    v_max_power: float = 77.0 # Max change for P_G, P_SC, P_SD per iteration
    v_max_theta: float = 0.9  # Max change for theta per iteration
    v_max_binary: float = 8.0 # Max velocity value before sigmoid for binary

# --- Helper Functions ---

def dominates(obj1, obj2):
    """Checks if obj1 dominates obj2."""
    # Assumes minimization for both objectives
    return np.all(obj1 <= obj2) and np.any(obj1 < obj2)

def sigmoid(x):
    """Sigmoid function for binary variable updates."""
    try:
        return 1 / (1 + math.exp(-x))
    except OverflowError:
        # Handle extreme values if math.exp overflows
        return 0.0 if x < 0 else 1.0

def calculate_detailed_costs(state: State, env: PowerSystemEnv):
    """Calculates the breakdown of costs for a given state."""
    sys_config = env.sys_config
    gen_config = env.gen_config
    storage_config = env.storage_config

    details = {}

    # Generation Costs
    gen_cost1 = sum((gen_config.a * state.P_G1 ** 2 + gen_config.b * state.P_G1 + gen_config.c) * sys_config.P_coal)
    gen_cost2 = sum((gen_config.a * state.P_G2 ** 2 + gen_config.b * state.P_G2 + gen_config.c) * sys_config.P_coal)
    gen_cost3 = sum((gen_config.a * state.P_G3 ** 2 + gen_config.b * state.P_G3 + gen_config.c) * sys_config.P_coal)
    details['gen_cost'] = gen_cost1 + gen_cost2 + gen_cost3

    # Switching Costs
    switch_cost1 = gen_config.switch_cost * sum(abs(state.u_G1[1:] - state.u_G1[:-1]))
    switch_cost2 = gen_config.switch_cost * sum(abs(state.u_G2[1:] - state.u_G2[:-1]))
    switch_cost3 = gen_config.switch_cost * sum(abs(state.u_G3[1:] - state.u_G3[:-1]))
    details['switch_cost'] = switch_cost1 + switch_cost2 + switch_cost3

    # Storage Costs (Operation)
    details['storage_op_cost'] = storage_config.charge_cost * sum(state.P_SC + state.P_SD)

    # Penalty Costs
    details['penalty_Wcut'] = sys_config.m_Wcut * sum(state.P_Wcut)
    details['penalty_PVcut'] = sys_config.m_PVcut * sum(state.P_PVcut)
    details['penalty_Lcut'] = sys_config.m_Lcut * sum(state.P_Lcut)
    details['penalty_cost_total'] = details['penalty_Wcut'] + details['penalty_PVcut'] + details['penalty_Lcut']

    # Carbon Trading Costs
    carbon_trading_cost1 = sys_config.sigma_e * sum(sys_config.E_S * state.P_G1 - sys_config.lambda_c * state.P_G1)
    carbon_trading_cost2 = sys_config.sigma_e * sum(sys_config.E_S * state.P_G2 - sys_config.lambda_c * state.P_G2)
    carbon_trading_cost3 = sys_config.sigma_e * sum(sys_config.E_S * state.P_G3 - sys_config.lambda_c * state.P_G3)
    details['carbon_trading_cost'] = carbon_trading_cost1 + carbon_trading_cost2 + carbon_trading_cost3

    # Carbon Capture Costs
    carbon_capture1 = sys_config.sigma_s * sys_config.rho * sum(state.P_G1 * sys_config.E_S * state.theta * sys_config.eta_b)
    carbon_capture2 = sys_config.sigma_s * sys_config.rho * sum(state.P_G2 * sys_config.E_S * state.theta * sys_config.eta_b)
    carbon_capture3 = sys_config.sigma_s * sys_config.rho * sum(state.P_G3 * sys_config.E_S * state.theta * sys_config.eta_b)
    details['carbon_capture_cost'] = carbon_capture1 + carbon_capture2 + carbon_capture3

    # Carbon Fen Costs (Penalty/Tax related to capture operation?)
    carbon_fen1 = sys_config.sigma_e * sys_config.K_S * sum(state.P_G1 * sys_config.E_S * state.theta * sys_config.eta_b)
    carbon_fen2 = sys_config.sigma_e * sys_config.K_S * sum(state.P_G2 * sys_config.E_S * state.theta * sys_config.eta_b)
    carbon_fen3 = sys_config.sigma_e * sys_config.K_S * sum(state.P_G3 * sys_config.E_S * state.theta * sys_config.eta_b)
    details['carbon_fen_cost'] = carbon_fen1 + carbon_fen2 + carbon_fen3

    # Total Cost (Recalculate for verification, should match objective 1)
    details['total_cost_recalc'] = (details['gen_cost'] + details['switch_cost'] +
                                   details['storage_op_cost'] + details['penalty_cost_total'] +
                                   details['carbon_capture_cost'] + details['carbon_trading_cost'] +
                                   details['carbon_fen_cost'])

    # Carbon Emissions (Recalculate for verification, should match objective 2)
    carbon_emission1 = sum(sys_config.E_S * state.P_G1 - state.P_G1 * sys_config.E_S * state.theta * sys_config.eta_b)
    carbon_emission2 = sum(sys_config.E_S * state.P_G2 - state.P_G2 * sys_config.E_S * state.theta * sys_config.eta_b)
    carbon_emission3 = sum(sys_config.E_S * state.P_G3 - state.P_G3 * sys_config.E_S * state.theta * sys_config.eta_b)
    details['carbon_emission_recalc'] = carbon_emission1 + carbon_emission2 + carbon_emission3

    return details

# --- MOPSO Core Logic ---

class Particle:
    def __init__(self, env: PowerSystemEnv, mopsoconfig: MOPSOConfig):
        self.env = env
        self.mopsoconfig = mopsoconfig
        self.N_T = env.N_T

        # Initialize Position
        self.position = env.create_random_state()
        # Important: Repair initial random state to satisfy constraints
        self.position = env.repair_solution(self.position)
        self.objectives = env.calculate_objectives(self.position)

        # Initialize Velocity (same structure as State, filled with zeros initially)
        self.velocity = State(
            P_G1=np.zeros(self.N_T), P_G2=np.zeros(self.N_T), P_G3=np.zeros(self.N_T),
            u_G1=np.zeros(self.N_T), u_G2=np.zeros(self.N_T), u_G3=np.zeros(self.N_T),
            theta=np.zeros(self.N_T),
            P_SC=np.zeros(self.N_T), P_SD=np.zeros(self.N_T),
            u_SC=np.zeros(self.N_T), u_SD=np.zeros(self.N_T),
            E_ES=np.zeros(self.N_T + 1), # Velocity for E_ES is not directly used
            P_Wcut=np.zeros(self.N_T), P_PVcut=np.zeros(self.N_T), P_Lcut=np.zeros(self.N_T) # Not directly used
        )
        # Initialize random velocities for relevant parts
        self._randomize_velocity()


        # Initialize Personal Best
        self.pbest_position = copy.deepcopy(self.position)
        self.pbest_objectives = self.objectives

    def _randomize_velocity(self):
        """Initializes velocities randomly within bounds."""
        max_power_range = self.env.gen_config.max_output - self.env.gen_config.min_output
        self.velocity.P_G1 = np.random.uniform(-max_power_range*0.1, max_power_range*0.1, self.N_T)
        self.velocity.P_G2 = np.random.uniform(-max_power_range*0.1, max_power_range*0.1, self.N_T)
        self.velocity.P_G3 = np.random.uniform(-max_power_range*0.1, max_power_range*0.1, self.N_T)

        self.velocity.theta = np.random.uniform(-0.1, 0.1, self.N_T) # Theta range [0,1]

        max_charge_range = self.env.storage_config.max_charge - self.env.storage_config.min_charge
        max_discharge_range = self.env.storage_config.max_discharge - self.env.storage_config.min_discharge
        self.velocity.P_SC = np.random.uniform(-max_charge_range*0.1, max_charge_range*0.1, self.N_T)
        self.velocity.P_SD = np.random.uniform(-max_discharge_range*0.1, max_discharge_range*0.1, self.N_T)

        # For binary variables, initialize velocity in a range suitable for sigmoid
        self.velocity.u_G1 = np.random.uniform(-self.mopsoconfig.v_max_binary, self.mopsoconfig.v_max_binary, self.N_T)
        self.velocity.u_G2 = np.random.uniform(-self.mopsoconfig.v_max_binary, self.mopsoconfig.v_max_binary, self.N_T)
        self.velocity.u_G3 = np.random.uniform(-self.mopsoconfig.v_max_binary, self.mopsoconfig.v_max_binary, self.N_T)
        self.velocity.u_SC = np.random.uniform(-self.mopsoconfig.v_max_binary, self.mopsoconfig.v_max_binary, self.N_T)
        self.velocity.u_SD = np.random.uniform(-self.mopsoconfig.v_max_binary, self.mopsoconfig.v_max_binary, self.N_T)


    def update_velocity(self, gbest_position: State, w: float):
        """Updates particle velocity."""
        c1 = self.mopsoconfig.c1
        c2 = self.mopsoconfig.c2
        r1 = random.random()
        r2 = random.random()

        v_max_p = self.mopsoconfig.v_max_power
        v_max_t = self.mopsoconfig.v_max_theta
        v_max_b = self.mopsoconfig.v_max_binary

        # --- Update Continuous Variables Velocity ---
        self.velocity.P_G1 = np.clip(w * self.velocity.P_G1 +
                                   c1 * r1 * (self.pbest_position.P_G1 - self.position.P_G1) +
                                   c2 * r2 * (gbest_position.P_G1 - self.position.P_G1), -v_max_p, v_max_p)
        self.velocity.P_G2 = np.clip(w * self.velocity.P_G2 +
                                   c1 * r1 * (self.pbest_position.P_G2 - self.position.P_G2) +
                                   c2 * r2 * (gbest_position.P_G2 - self.position.P_G2), -v_max_p, v_max_p)
        self.velocity.P_G3 = np.clip(w * self.velocity.P_G3 +
                                   c1 * r1 * (self.pbest_position.P_G3 - self.position.P_G3) +
                                   c2 * r2 * (gbest_position.P_G3 - self.position.P_G3), -v_max_p, v_max_p)

        self.velocity.theta = np.clip(w * self.velocity.theta +
                                    c1 * r1 * (self.pbest_position.theta - self.position.theta) +
                                    c2 * r2 * (gbest_position.theta - self.position.theta), -v_max_t, v_max_t)

        self.velocity.P_SC = np.clip(w * self.velocity.P_SC +
                                   c1 * r1 * (self.pbest_position.P_SC - self.position.P_SC) +
                                   c2 * r2 * (gbest_position.P_SC - self.position.P_SC), -v_max_p, v_max_p)
        self.velocity.P_SD = np.clip(w * self.velocity.P_SD +
                                   c1 * r1 * (self.pbest_position.P_SD - self.position.P_SD) +
                                   c2 * r2 * (gbest_position.P_SD - self.position.P_SD), -v_max_p, v_max_p)

        # --- Update Binary Variables Velocity ---
        # (Note: pbest/gbest positions are 0/1, direct subtraction works)
        # Convert binary arrays to integers before subtraction
        self.velocity.u_G1 = np.clip(w * self.velocity.u_G1 +
                                   c1 * r1 * (self.pbest_position.u_G1.astype(int) - self.position.u_G1.astype(int)) +
                                   c2 * r2 * (gbest_position.u_G1.astype(int) - self.position.u_G1.astype(int)), -v_max_b, v_max_b)
        self.velocity.u_G2 = np.clip(w * self.velocity.u_G2 +
                                   c1 * r1 * (self.pbest_position.u_G2.astype(int) - self.position.u_G2.astype(int)) +
                                   c2 * r2 * (gbest_position.u_G2.astype(int) - self.position.u_G2.astype(int)), -v_max_b, v_max_b)
        self.velocity.u_G3 = np.clip(w * self.velocity.u_G3 +
                                   c1 * r1 * (self.pbest_position.u_G3.astype(int) - self.position.u_G3.astype(int)) +
                                   c2 * r2 * (gbest_position.u_G3.astype(int) - self.position.u_G3.astype(int)), -v_max_b, v_max_b)
        self.velocity.u_SC = np.clip(w * self.velocity.u_SC +
                                   c1 * r1 * (self.pbest_position.u_SC.astype(int) - self.position.u_SC.astype(int)) +
                                   c2 * r2 * (gbest_position.u_SC.astype(int) - self.position.u_SC.astype(int)), -v_max_b, v_max_b)
        self.velocity.u_SD = np.clip(w * self.velocity.u_SD +
                                   c1 * r1 * (self.pbest_position.u_SD.astype(int) - self.position.u_SD.astype(int)) +
                                   c2 * r2 * (gbest_position.u_SD.astype(int) - self.position.u_SD.astype(int)), -v_max_b, v_max_b)


    def update_position(self):
        """Updates particle position based on velocity."""
        new_pos = copy.deepcopy(self.position) # Start with current position

        # --- Update Continuous Variables ---
        new_pos.P_G1 = self.position.P_G1 + self.velocity.P_G1
        new_pos.P_G2 = self.position.P_G2 + self.velocity.P_G2
        new_pos.P_G3 = self.position.P_G3 + self.velocity.P_G3
        new_pos.theta = self.position.theta + self.velocity.theta
        new_pos.P_SC = self.position.P_SC + self.velocity.P_SC
        new_pos.P_SD = self.position.P_SD + self.velocity.P_SD

        # --- Update Binary Variables using Sigmoid ---
        rand_vals = np.random.rand(self.N_T) # Generate random numbers once per variable
        new_pos.u_G1 = (rand_vals < np.array([sigmoid(v) for v in self.velocity.u_G1])).astype(int)

        rand_vals = np.random.rand(self.N_T)
        new_pos.u_G2 = (rand_vals < np.array([sigmoid(v) for v in self.velocity.u_G2])).astype(int)

        rand_vals = np.random.rand(self.N_T)
        new_pos.u_G3 = (rand_vals < np.array([sigmoid(v) for v in self.velocity.u_G3])).astype(int)

        rand_vals = np.random.rand(self.N_T)
        new_pos.u_SC = (rand_vals < np.array([sigmoid(v) for v in self.velocity.u_SC])).astype(int)

        # Ensure u_SD is opposite of u_SC if both happen to be 1 after update
        # (Repair step will handle final consistency)
        rand_vals = np.random.rand(self.N_T)
        new_pos.u_SD = (rand_vals < np.array([sigmoid(v) for v in self.velocity.u_SD])).astype(int)
        # Simple check: if u_SC is 1, force u_SD to 0 for now
        new_pos.u_SD[new_pos.u_SC == 1] = 0


        # --- Repair and Evaluate ---
        # Crucial step: Ensure the new position is feasible
        repaired_pos = self.env.repair_solution(new_pos)

        # Update particle's actual position and objectives
        self.position = repaired_pos
        self.objectives = self.env.calculate_objectives(self.position)


    def update_pbest(self):
        """Updates the particle's personal best."""
        new_obj = self.objectives
        pbest_obj = self.pbest_objectives

        if dominates(new_obj, pbest_obj):
            self.pbest_position = copy.deepcopy(self.position)
            self.pbest_objectives = new_obj
        # Optional: If they are non-dominated, could potentially add to a pbest archive
        # For simplicity here, we only replace if strictly dominating.


class MOPSO:
    def __init__(self, env: PowerSystemEnv, mopsoconfig: MOPSOConfig):
        self.env = env
        self.mopsoconfig = mopsoconfig
        self.pop_size = mopsoconfig.pop_size
        self.n_generations = mopsoconfig.n_generations
        self.archive = [] # List of {'position': State, 'objectives': np.array}
        self.particles = [Particle(env, mopsoconfig) for _ in range(self.pop_size)]
        self._initialize_archive()

    def _initialize_archive(self):
        """Initializes the archive with non-dominated solutions from the initial population."""
        print("Initializing archive...")
        initial_solutions = [(p.position, p.objectives) for p in self.particles]
        for pos, obj in initial_solutions:
            self._add_to_archive(pos, obj)
        print(f"Initial archive size: {len(self.archive)}")


    def _add_to_archive(self, position: State, objectives: np.ndarray):
        """Adds a solution to the archive, maintaining non-domination."""
        new_obj = np.array(objectives)

        # Check if the new solution is dominated by any existing archive member
        if any(dominates(existing['objectives'], new_obj) for existing in self.archive):
            return # New solution is dominated, do not add

        # Remove archive members that are dominated by the new solution
        self.archive = [entry for entry in self.archive if not dominates(new_obj, entry['objectives'])]

        # Add the new non-dominated solution
        self.archive.append({'position': copy.deepcopy(position), 'objectives': new_obj})

        # Optional: Prune archive if it exceeds the size limit
        if self.mopsoconfig.archive_size_limit > 0 and len(self.archive) > self.mopsoconfig.archive_size_limit:
            self._prune_archive()

    def _prune_archive(self):
         # Simple pruning: randomly remove solutions until size limit is met
         # More sophisticated methods (crowding distance) could be used here
         while len(self.archive) > self.mopsoconfig.archive_size_limit:
             remove_idx = random.randrange(len(self.archive))
             self.archive.pop(remove_idx)
         # print(f"Archive pruned to size: {len(self.archive)}") # Debug


    def _select_leader(self) -> State:
        """Selects a leader (gbest) from the archive."""
        if not self.archive:
             # Should not happen after initialization, but fallback if it does
             return self.particles[random.randrange(self.pop_size)].pbest_position
        # Simple random selection from archive
        return random.choice(self.archive)['position']

    def run(self):
        """Executes the MOPSO algorithm."""
        print("Starting MOPSO optimization...")
        for gen in tqdm(range(self.n_generations), desc="Generations"):
            # Update inertia weight linearly
            w = self.mopsoconfig.w_max - (self.mopsoconfig.w_max - self.mopsoconfig.w_min) * (gen / self.n_generations)

            for particle in self.particles:
                # Select leader for this particle
                gbest_position = self._select_leader()

                # Update velocity and position
                particle.update_velocity(gbest_position, w)
                particle.update_position() # Includes repair and evaluation

                # Update personal best
                particle.update_pbest()

                # Try adding the particle's new position to the archive
                self._add_to_archive(particle.position, particle.objectives)

            # Optional: Print archive size periodically
            # if (gen + 1) % 10 == 0:
            #     print(f"Generation {gen + 1}/{self.n_generations}, Archive size: {len(self.archive)}")

        print(f"Optimization finished. Final archive size: {len(self.archive)}")
        return self.archive

    def plot_pareto_front(self, filename="pareto_front_mopso.png"):
        """Plots the obtained Pareto front and saves data."""
        if not self.archive:
            print("Archive is empty, cannot plot.")
            return
    
        objectives = np.array([entry['objectives'] for entry in self.archive])
        costs = objectives[:, 0]
        emissions = objectives[:, 1]
    
        # Save data to CSV
        data_path = os.path.join('data', 'pareto_points.csv')
        df = pd.DataFrame({
            'Operation_Cost': costs,
            'Carbon_Emissions': emissions
        })
        df.to_csv(data_path, index=False)
        print(f"Pareto points data saved to {data_path}")
    
        # Original plotting code continues...
        plt.figure(figsize=(10, 6))
        plt.scatter(costs, emissions, c='blue', marker='o', label='Pareto Front Solutions')
        plt.title('Pareto Front (MOPSO)')
        plt.xlabel('Total Operating Cost ($)')
        plt.ylabel('Total Carbon Emissions (Units)')
        plt.grid(True)
        plt.legend()
        plt.savefig(filename)
        print(f"Pareto front plot saved to {filename}")
        plt.show()


    def save_detailed_results(self, filename="pareto_solutions_details_mopso.csv"):
        """Saves detailed cost/emission breakdown for Pareto solutions to CSV."""
        if not self.archive:
            print("Archive is empty, cannot save detailed results.")
            return

        results_data = []
        solution_index = 0
        for entry in self.archive:
            solution_index += 1
            state = entry['position']
            total_cost, total_emission = entry['objectives']
            detailed_costs = calculate_detailed_costs(state, self.env)

            row = {
                'Solution': solution_index,
                'Total Cost': total_cost,
                'Total Emission': total_emission,
                **detailed_costs # Add all detailed costs
                # Optionally add state variables if needed, but makes file huge
                # 'P_G1': state.P_G1.tolist(), # Example
            }
            results_data.append(row)

        results_df = pd.DataFrame(results_data)
        # Optional: Verify totals match
        cost_match = np.allclose(results_df['Total Cost'], results_df['total_cost_recalc'])
        emission_match = np.allclose(results_df['Total Emission'], results_df['carbon_emission_recalc'])
        print(f"Verification: Total Cost matches detailed sum: {cost_match}")
        print(f"Verification: Total Emission matches detailed sum: {emission_match}")

        results_df = results_df.drop(columns=['total_cost_recalc', 'carbon_emission_recalc']) # Remove verification columns

        results_df.to_csv(filename, index=False)
        print(f"Detailed results for Pareto solutions saved to {filename}")


# --- Main Execution ---
if __name__ == "__main__":
    # 1. Load configurations
    # Make sure 'data/forecasts.csv' exists or change the path
    try:
        system_config, generator_config, storage_config, _, load_f, wind_f, pv_f = get_default_configs(data_filepath='data/forecasts.csv')
        print("Configurations and forecast data loaded successfully.")
    except FileNotFoundError:
        print("Error: data/forecasts.csv not found. Please ensure the file exists.")
        # Create dummy data if file not found, for testing purposes ONLY
        print("Using dummy forecast data for testing.")
        N_T = 24
        system_config = SystemConfig(N_T=N_T)
        generator_config = GeneratorConfig()
        storage_config = StorageConfig()
        load_f = np.random.uniform(400, 800, N_T)
        wind_f = np.random.uniform(0, 200, N_T)
        pv_f = np.random.uniform(0, 150, N_T)


    # 2. Initialize Environment
    env = PowerSystemEnv(system_config, generator_config, storage_config,
                         load_f, wind_f, pv_f)
    print("PowerSystemEnv initialized.")

    # 3. Setup MOPSO
    mopsoconfig = MOPSOConfig() # Use default MOPSO settings
    mopso_solver = MOPSO(env, mopsoconfig)

    # 4. Run Optimization
    final_archive = mopso_solver.run()

    # 5. Process Results
    if final_archive:
        mopso_solver.plot_pareto_front()
        mopso_solver.save_detailed_results()
    else:
        print("MOPSO did not find any non-dominated solutions.")